# -*- coding: utf-8 -*-
"""
阿里云语音识别客户端
基于阿里云官方WebSocket API实现
参考文档: https://help.aliyun.com/zh/model-studio/websocket-for-paraformer-real-time-service
"""

import asyncio
import json
import logging
import uuid
import websockets
from typing import Optional, Callable, Dict, Any
from settings import settings

logger = logging.getLogger(__name__)

class AliyunASRClient:
    """阿里云语音识别客户端 - 基于官方WebSocket API"""
    
    def __init__(self, 
                 model: str = None,
                 format: str = 'wav',
                 sample_rate: int = 16000,
                 language_hints: list = None,
                 on_result_callback: Optional[Callable[[str], None]] = None):
        """
        初始化阿里云ASR客户端
        
        Args:
            model: 模型名称，默认使用settings中的配置
            format: 音频格式，默认wav
            sample_rate: 采样率，默认16000
            language_hints: 语言提示，仅支持paraformer-realtime-v2模型
            on_result_callback: 识别结果回调函数
        """
        self.model = model or settings.ALIYUN_MODEL
        self.format = format
        self.sample_rate = sample_rate
        self.language_hints = language_hints or ['zh', 'en']
        self.on_result_callback = on_result_callback
        
        # WebSocket连接
        self.websocket = None
        self.is_connected = False
        self.task_id = str(uuid.uuid4())
        
        # 任务状态
        self.task_started = False  # 是否收到task-started事件
        
        # 音频缓冲区
        self.audio_buffer = bytearray()
        
        # 消息监听任务
        self.listen_task = None
        
        # 设置API Key
        if not settings.ALIYUN_API_KEY:
            logger.warning("阿里云API Key未配置，请设置ALIYUN_API_KEY环境变量")
    
    async def connect(self):
        """连接阿里云WebSocket服务（官方API方式）"""
        try:
            logger.info(f"[阿里云ASR] 正在连接，模型: {self.model}")
            # 官方WebSocket URL
            ws_url = "wss://dashscope.aliyuncs.com/api-ws/v1/inference"
            # 官方要求的Header
            headers = {
                "Authorization": f"bearer {settings.ALIYUN_API_KEY}",
                "user-agent": "fastapi-ota-demo"
            }
            self.websocket = await websockets.connect(ws_url, additional_headers=headers)
            self.is_connected = True
            logger.info(f"[阿里云ASR] 连接成功")
        except Exception as e:
            logger.error(f"[阿里云ASR] 连接失败: {e}")
            self.is_connected = False
            raise
    
    async def start(self):
        """开始识别会话"""
        if not self.is_connected or not self.websocket:
            logger.error("[阿里云ASR] 客户端未连接")
            return False
        
        try:
            logger.info("[阿里云ASR] 开始识别任务")
            
            # 构建run-task指令，参考官方文档
            run_task_message = {
                "header": {
                    "action": "run-task",
                    "task_id": self.task_id,
                    "streaming": "duplex"
                },
                "payload": {
                    "task_group": "audio",
                    "task": "asr",
                    "function": "recognition",
                    "model": self.model,
                    "parameters": {
                        "sample_rate": self.sample_rate,
                        "format": self.format
                    },
                    "input": {}
                }
            }
            
            # 添加语言提示（仅支持paraformer-realtime-v2）
            if self.model == "paraformer-realtime-v2" and self.language_hints:
                run_task_message["payload"]["parameters"]["language_hints"] = self.language_hints
            
            # 发送run-task指令
            await self.websocket.send(json.dumps(run_task_message))
            
            # 启动消息监听任务
            self.listen_task = asyncio.create_task(self._listen_messages())
            
            return True
            
        except Exception as e:
            logger.error(f"[阿里云ASR] 开始识别失败: {e}")
            return False
    
    async def _listen_messages(self):
        """监听WebSocket消息"""
        try:
            async for message in self.websocket:
                try:
                    # 解析消息
                    if isinstance(message, str):
                        # JSON消息
                        data = json.loads(message)
                        await self._handle_json_message(data)
                    else:
                        # 二进制消息（音频数据）
                        await self._handle_binary_message(message)
                        
                except json.JSONDecodeError as e:
                    logger.error(f"[阿里云ASR] JSON解析失败: {e}")
                except Exception as e:
                    logger.error(f"[阿里云ASR] 处理消息失败: {e}")
                    
        except websockets.exceptions.ConnectionClosed:
            logger.info("[阿里云ASR] WebSocket连接已关闭")
        except Exception as e:
            logger.error(f"[阿里云ASR] 监听消息失败: {e}")
        finally:
            self.is_connected = False
    
    async def _handle_json_message(self, data: Dict):
        """处理JSON消息"""
        try:
            header = data.get("header", {})
            event = header.get("event")  # 事件类型

            if event == "task-started":
                self.task_started = True
                logger.info("[阿里云ASR] 识别任务已开始，可以发送音频数据")

            elif event == "result-generated":
                payload = data.get("payload", {})
                output = payload.get("output", {})
                sentence = output.get("sentence", {})

                text = sentence.get("text", "")
                sentence_end = sentence.get("sentence_end", False)
                if text and text.strip():
                    logger.info(f"[阿里云ASR] 识别结果: {text} (sentence_end={sentence_end})")
                    # 回调函数支持新参数
                    if self.on_result_callback:
                        try:
                            if asyncio.iscoroutinefunction(self.on_result_callback):
                                # 新版：支持(text, sentence_end)
                                await self.on_result_callback(text, sentence_end)
                            else:
                                # 检查回调参数个数，兼容老用法
                                import inspect
                                params = inspect.signature(self.on_result_callback).parameters
                                if len(params) == 2:
                                    self.on_result_callback(text, sentence_end)
                                else:
                                    self.on_result_callback(text)
                        except Exception as e:
                            logger.error(f"[阿里云ASR] 回调函数调用出错: {e}")

                if sentence_end:
                    logger.info("[阿里云ASR] 句子识别完成")

            elif event == "task-finished":
                logger.info("[阿里云ASR] 识别任务完成")

            elif event == "task-failed":
                error_message = header.get("error_message", "未知错误")
                logger.error(f"[阿里云ASR] 任务失败: {error_message}")

            else:
                logger.debug(f"[阿里云ASR] 收到事件: {event}")

        except Exception as e:
            logger.error(f"[阿里云ASR] 处理JSON消息失败: {e}")
    
    async def _handle_binary_message(self, data: bytes):
        """处理二进制消息（音频数据）"""
        # 阿里云WebSocket API通常不返回二进制音频数据
        logger.debug(f"[阿里云ASR] 收到二进制数据: {len(data)} 字节")
    
    async def send_audio(self, audio_chunk: bytes, is_last: bool = False):
        """
        发送音频数据
        
        Args:
            audio_chunk: 音频数据块
            is_last: 是否为最后一帧
        """
        if not self.is_connected or not self.websocket:
            logger.error("[阿里云ASR] 客户端未连接")
            return
        
        if not self.task_started:
            logger.warning("[阿里云ASR] 任务未开始，等待task-started事件")
            return
        
        try:
            # logger.info(f"[阿里云ASR] 收到音频数据块，长度: {len(audio_chunk)} 字节")
            
            # 将音频数据添加到缓冲区
            self.audio_buffer.extend(audio_chunk)
            # logger.info(f"[阿里云ASR] 缓冲区当前长度: {len(self.audio_buffer)} 字节")
            
            # 按100ms分包发送（16k 16bit 单声道 = 3200字节/100ms）
            chunk_size = 1600 if self.sample_rate == 8000 else 3200  # 16k采样率
            
            while len(self.audio_buffer) >= chunk_size:
                chunk = bytes(self.audio_buffer[:chunk_size])
                self.audio_buffer = self.audio_buffer[chunk_size:]
                
                # 根据阿里云官方文档，音频数据应该直接作为二进制数据发送
                # 而不是封装在JSON中
                # logger.info(f"[阿里云ASR] 发送音频包到阿里云，长度: {len(chunk)} 字节")
                await self.websocket.send(chunk)
                await asyncio.sleep(0.1)  # 100ms间隔，参考官方文档
            
            # 如果是最后一帧，发送剩余数据
            if is_last and len(self.audio_buffer) > 0:
                logger.info(f"[阿里云ASR] 发送最后一帧音频数据，长度: {len(self.audio_buffer)} 字节")
                await self.websocket.send(bytes(self.audio_buffer))
                self.audio_buffer.clear()
                
        except Exception as e:
            logger.error(f"[阿里云ASR] 发送音频数据失败: {e}")
            # 添加更详细的错误信息
            logger.error(f"[阿里云ASR] 音频数据长度: {len(audio_chunk)}")
            logger.error(f"[阿里云ASR] 缓冲区长度: {len(self.audio_buffer)}")
            logger.error(f"[阿里云ASR] WebSocket连接状态: {self.is_connected}")
            logger.error(f"[阿里云ASR] 任务状态: {self.task_started}")
    
    async def stop(self):
        """停止识别会话"""
        if not self.websocket:
            logger.warning("[阿里云ASR] WebSocket连接不存在")
            return
        
        try:
            logger.info("[阿里云ASR] 结束识别任务")
            
            # 构建finish-task指令，参考官方文档
            finish_task_message = {
                "header": {
                    "action": "finish-task",
                    "task_id": self.task_id,
                    "streaming": "duplex"
                },
                "payload": {
                    "input": {}
                }
            }
            
            # 发送finish-task指令
            await self.websocket.send(json.dumps(finish_task_message))
            
            # 等待一段时间让服务器处理完成
            await asyncio.sleep(0.5)
            
        except Exception as e:
            logger.error(f"[阿里云ASR] 结束任务失败: {e}")
    
    async def close(self):
        """关闭连接"""
        try:
            # 取消监听任务
            if self.listen_task and not self.listen_task.done():
                self.listen_task.cancel()
                try:
                    await self.listen_task
                except asyncio.CancelledError:
                    pass
            
            # 关闭WebSocket连接
            if self.websocket:
                await self.websocket.close()
            
            self.is_connected = False
            self.audio_buffer.clear()
            logger.info("[阿里云ASR] 连接已关闭")
            
        except Exception as e:
            logger.error(f"[阿里云ASR] 关闭连接失败: {e}")
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取性能指标"""
        return {
            "task_id": self.task_id,
            "is_connected": self.is_connected,
            "model": self.model,
            "sample_rate": self.sample_rate
        }


class AliyunASRManager:
    """阿里云ASR管理器，用于管理多个客户端连接"""
    
    def __init__(self):
        self.clients: Dict[str, AliyunASRClient] = {}
    
    async def create_client(self, client_id: str, **kwargs) -> AliyunASRClient:
        """创建新的ASR客户端"""
        if client_id in self.clients:
            await self.remove_client(client_id)
        
        client = AliyunASRClient(**kwargs)
        self.clients[client_id] = client
        return client
    
    async def get_client(self, client_id: str) -> Optional[AliyunASRClient]:
        """获取客户端"""
        return self.clients.get(client_id)
    
    async def remove_client(self, client_id: str):
        """移除客户端"""
        if client_id in self.clients:
            client = self.clients[client_id]
            await client.close()
            del self.clients[client_id]
    
    async def close_all(self):
        """关闭所有客户端"""
        for client_id in list(self.clients.keys()):
            await self.remove_client(client_id)


# 全局ASR管理器实例
asr_manager = AliyunASRManager() 